"""
实现计算图和profiler结果的匹配，profiler获取到的数据全部都存储在计算图的meta字典中
"""
import torch
import torch.nn as nn
import torch.fx
import graphviz
import pickle
from optype_torch_map import pytorch_to_op_type
import json
from PIL import Image
from parse_config import Config
import os
import re
import subprocess

class LayerNode:
    def __init__(self, name):
        self.name = name
        self.input_shape = []
        self.output_shape = []
        self.input_memory = 0
        self.time_elapsed = 0
        self.peak_memory = 0

    def __repr__(self):
        return (f"LayerNode(name={self.name}, input_shape={self.input_shape}, "
                f"output_shape={self.output_shape}, input_memory={self.input_memory:.2f} MB, "
                f"time_elapsed={self.time_elapsed:.6f} s, peak_memory={self.peak_memory:.2f} MB)")


def get_image_height(image_path):
    """通过 Pillow 获取图片高度（单位：像素）"""
    with Image.open(image_path) as img:
        return img.height


def get_compute_nodes(pth_path,profile_nodes,meta_save_path,plain_nodes_save_path,computation_pic_save_path):
    # 从文件中加载计算图
    loaded_traced = torch.load(pth_path)

    # 创建一个 Graphviz 的 Digraph 对象
    dot = graphviz.Digraph(comment='ResNet18 Computation Graph')
    
    dpi = 300
    dot.attr(dpi=str(dpi))
    
    i=0

    # 遍历计算图中的节点
    for node in loaded_traced.graph.nodes:
        if node.op == 'call_module':
            # 获取被调用的模块
            module = loaded_traced.get_submodule(node.target)
            op_type = type(module).__name__
        else:
            op_type = node.op
        if i<len(profile_nodes) and op_type == profile_nodes[i].name:
            node.meta['input_shape']=profile_nodes[i].input_shape
            node.meta['output_shape']=profile_nodes[i].output_shape
            node.meta['input_memory']=profile_nodes[i].input_memory
            node.meta['time_elapsed']=profile_nodes[i].time_elapsed
            node.meta['peak_memory']=profile_nodes[i].peak_memory
            node.meta['output_memory']=profile_nodes[i].output_memory
            i+=1
        else:
            node.meta['input_shape']=[]
            node.meta['output_shape']=[]
            node.meta['input_memory']=0
            node.meta['time_elapsed']=0
            node.meta['peak_memory']=0
            if i+1 < len(profile_nodes):
                node.meta['output_memory']=profile_nodes[i+1].input_memory
            else:
                node.meta['output_memory']=0
        
        # 添加节点到图中
        dot.node(node.name, str(node))
        # 遍历节点的输入
        for input_node in node.all_input_nodes:
            # 添加边到图中
            dot.edge(input_node.name, node.name)
            
    meta_data=[]
    for node in loaded_traced.graph.nodes:
        meta={
            'input_shape':node.meta['input_shape'],
            'output_shape':node.meta['output_shape'],
            'input_memory':node.meta['input_memory'],
            'time_elapsed':node.meta['time_elapsed'],
            'peak_memory':node.meta['peak_memory'],
            'output_memory':node.meta['output_memory']
        }
        meta_data.append(meta)

    
    with open(meta_save_path,'w') as f:
        json.dump(meta_data, f)

    plain_output = dot.pipe(format='plain').decode('utf-8')
    
    
    dot_file_path = os.path.splitext(computation_pic_save_path)[0] + ".dot"
    dot.save(dot_file_path)  # 保存为 DOT 文件
    print(f"DOT file saved to: {dot_file_path}")
    
    error=''
    try:
        # 要运行的 Python 脚本
        script_path = 'load_dot.py'
        # 运行脚本并捕获输出
        result = subprocess.run(['python', script_path], capture_output=True, text=True, check=True)
        error=result.stderr
    except subprocess.CalledProcessError as e:
        print(f"运行脚本时出错: {e}")
        
    scale=1
    match = re.search(r'Scaling by (\d+\.\d+) to fit', error)
    if match:
        scale = float(match.group(1))
    
    # 获取渲染后图片的高度（单位：像素）
    image_path = computation_pic_save_path
    image_height = get_image_height(image_path)
    
    # plain 输出中每一行以 "node", "edge" 或 "graph" 开头
    # 节点行格式为：
    # node <name> <x> <y> <width> <height> <label> <style> <shape> <color> <fillcolor>
    node_positions = {}
    for line in plain_output.splitlines():
        parts = line.split()
        if parts and parts[0] == 'node':
            name = parts[1]
            y_plain = float(parts[3]) * dpi
            y_img = image_height - y_plain*scale
            y_start=y_img-float(parts[5])*dpi/2*scale
            y_end=y_img+float(parts[5])*dpi/2*scale
            node_positions[name] = ((y_start-90*scale),(y_end+40*scale))
                
    with open(plain_nodes_save_path,'w') as f:
        json.dump(node_positions, f)

    
def get_profile(profile_path):
    with open(profile_path, 'rb') as f:
        loaded_layer_nodes = pickle.load(f)
    profile_nodes=[node for node in loaded_layer_nodes if node.name in pytorch_to_op_type.keys()]
    return profile_nodes
    
    
def main():
    config_path='./config.json'
    config=Config(config_path)
    profile_nodes= get_profile(config.config['profiler_result_save_path'])
    get_compute_nodes(config.config['formal_model_save_path'],profile_nodes,config.config['meta_result_save_path'],config.config['plain_result_save_path'],config.config['formal_computation_pic_save_path'])
    
if __name__== '__main__':
    main()